// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


// Basic memory arena.
//
// Standard usage:
//   char* pos = arena->reserve(100)
//   write data to [pos, pos + 100)
//   pos = arena->reserve(27);
//   write data to [pos, pos + 27)
//   and so on

#pragma once

#include <cstdint>
#include <cstring>
#include <memory>
#include <string_view>
#include <vector>

namespace kumo::pollux {
    // Internally manages memory in chunks. Releases memory only upon destruction.
    // Arena is NOT threadsafe: external locking is required. All functions in this
    // class are expected to be used in tight loops, so we inline everything.
    class Arena {
    public:
        Arena(int64_t initial_chunk_size = kMinChunkSize) {
            addChunk(initial_chunk_size);
            reserveEnd_ = pos_;
        }

        // Returns a pointer to a block of memory of size bytes that can be written
        // to, and guarantees for the lifetime of *this that that region will remain
        // valid. Does NOT guarantee that the region is initially 0'd.
        char *reserve(int64_t bytes) {
            if (reserveEnd_ + bytes <= chunkEnd_) {
                pos_ = reserveEnd_;
                reserveEnd_ += bytes;
            } else {
                addChunk(bytes);
            }
            return pos_;
        }

        // Copies |data| into the chunk, returning a view to the copied data.
        std::string_view writeString(std::string_view data) {
            char *pos = reserve(data.size());
            memcpy(pos, data.data(), data.size());
            return {pos, data.size()};
        }

    private:
        static constexpr int64_t kMinChunkSize = 1LL << 20;

        void addChunk(int64_t bytes) {
            const int64_t chunkSize = std::max(bytes, kMinChunkSize);
            chunks_.emplace_back(new char[chunkSize]);
            pos_ = chunks_.back().get();
            chunkEnd_ = pos_ + chunkSize;
            reserveEnd_ = pos_ + bytes;
        }

        char *chunkEnd_;
        char *pos_;
        char *reserveEnd_;
        std::vector<std::unique_ptr<char[]> > chunks_;
    };
} //  namespace kumo::pollux
